{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Diffusion-based Policy Learning for RL\n",
    "\n",
    "This notebook implements Diffusion Policy, a diffusion model that predicts robot action sequences in reinforcement learning tasks.\n",
    "\n",
    "This example implements a robot control model for pushing a T-shaped block into a target area. The model takes in current state observations as input, and outputs a trajectory of subsequent steps to follow. This script was contributed by [Dorsa Rohani](https://github.com/DorsaRoh) and the notebook by [Parag Ekbote](https://github.com/ParagEkbote)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in links: https://download.pytorch.org/whl/torch_stable.html\n",
      "Collecting git+https://github.com/rail-berkeley/d4rl.git\n",
      "  Cloning https://github.com/rail-berkeley/d4rl.git to /tmp/pip-req-build-tdkn3r22\n",
      "  Running command git clone --filter=blob:none --quiet https://github.com/rail-berkeley/d4rl.git /tmp/pip-req-build-tdkn3r22\n",
      "\n",
      "  Resolved https://github.com/rail-berkeley/d4rl.git to commit 89141a689b0353b0dac3da5cba60da4b1b16254d\n",
      "  Preparing metadata (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25hCollecting torch==2.0.1+cu117\n",
      "  Downloading https://download.pytorch.org/whl/cu117/torch-2.0.1%2Bcu117-cp310-cp310-linux_x86_64.whl (1843.9 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.8/1.8 GB\u001b[0m \u001b[31m49.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m:00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting torchvision==0.15.2+cu117\n",
      "  Downloading https://download.pytorch.org/whl/cu117/torchvision-0.15.2%2Bcu117-cp310-cp310-linux_x86_64.whl (6.1 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m6.1/6.1 MB\u001b[0m \u001b[31m192.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting torchaudio==2.0.2+cu117\n",
      "  Downloading https://download.pytorch.org/whl/cu117/torchaudio-2.0.2%2Bcu117-cp310-cp310-linux_x86_64.whl (4.4 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.4/4.4 MB\u001b[0m \u001b[31m181.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting gym==0.23.1\n",
      "  Downloading gym-0.23.1.tar.gz (626 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m626.2/626.2 kB\u001b[0m \u001b[31m99.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Installing build dependencies ... \u001b[?25ldone\n",
      "\u001b[?25h  Getting requirements to build wheel ... \u001b[?25ldone\n",
      "\u001b[?25h  Preparing metadata (pyproject.toml) ... \u001b[?25ldone\n",
      "\u001b[?25hCollecting protobuf==3.20.1\n",
      "  Downloading protobuf-3.20.1-cp310-cp310-manylinux_2_12_x86_64.manylinux2010_x86_64.whl.metadata (698 bytes)\n",
      "Collecting einops\n",
      "  Downloading einops-0.8.1-py3-none-any.whl.metadata (13 kB)\n",
      "Collecting mediapy\n",
      "  Downloading mediapy-1.2.2-py3-none-any.whl.metadata (4.8 kB)\n",
      "Collecting Pillow==9.0.0\n",
      "  Downloading Pillow-9.0.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (7.6 kB)\n",
      "Requirement already satisfied: filelock in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from torch==2.0.1+cu117) (3.17.0)\n",
      "Requirement already satisfied: typing-extensions in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from torch==2.0.1+cu117) (4.12.2)\n",
      "Requirement already satisfied: sympy in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from torch==2.0.1+cu117) (1.13.3)\n",
      "Requirement already satisfied: networkx in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from torch==2.0.1+cu117) (3.4.2)\n",
      "Requirement already satisfied: jinja2 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from torch==2.0.1+cu117) (3.1.5)\n",
      "Collecting triton==2.0.0 (from torch==2.0.1+cu117)\n",
      "  Downloading triton-2.0.0-1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.0 kB)\n",
      "Requirement already satisfied: numpy in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from torchvision==0.15.2+cu117) (1.26.4)\n",
      "Requirement already satisfied: requests in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from torchvision==0.15.2+cu117) (2.32.3)\n",
      "Requirement already satisfied: cloudpickle>=1.2.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from gym==0.23.1) (3.1.1)\n",
      "Collecting gym_notices>=0.0.4 (from gym==0.23.1)\n",
      "  Downloading gym_notices-0.0.8-py3-none-any.whl.metadata (1.0 kB)\n",
      "Collecting cmake (from triton==2.0.0->torch==2.0.1+cu117)\n",
      "  Downloading cmake-3.31.4-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (6.5 kB)\n",
      "Collecting lit (from triton==2.0.0->torch==2.0.1+cu117)\n",
      "  Downloading lit-18.1.8-py3-none-any.whl.metadata (2.5 kB)\n",
      "Collecting mjrl@ git+https://github.com/aravindr93/mjrl@master#egg=mjrl (from D4RL==1.1)\n",
      "  Cloning https://github.com/aravindr93/mjrl (to revision master) to /tmp/pip-install-sr3n9qkg/mjrl_98dfd6b68c6048399cc4826c3796d7e4\n",
      "  Running command git clone --filter=blob:none --quiet https://github.com/aravindr93/mjrl /tmp/pip-install-sr3n9qkg/mjrl_98dfd6b68c6048399cc4826c3796d7e4\n",
      "\n",
      "  Resolved https://github.com/aravindr93/mjrl to commit 3871d93763d3b49c4741e6daeaebbc605fe140dc\n",
      "  Preparing metadata (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25hCollecting mujoco_py (from D4RL==1.1)\n",
      "  Downloading mujoco_py-2.1.2.14-py3-none-any.whl.metadata (669 bytes)\n",
      "Collecting pybullet (from D4RL==1.1)\n",
      "  Downloading pybullet-3.2.7-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (1.8 kB)\n",
      "Collecting h5py (from D4RL==1.1)\n",
      "  Downloading h5py-3.12.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (2.5 kB)\n",
      "Requirement already satisfied: termcolor in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from D4RL==1.1) (2.5.0)\n",
      "Requirement already satisfied: click in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from D4RL==1.1) (8.1.8)\n",
      "Collecting dm_control>=1.0.3 (from D4RL==1.1)\n",
      "  Downloading dm_control-1.0.27-py3-none-any.whl.metadata (1.6 kB)\n",
      "Requirement already satisfied: ipython in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from mediapy) (8.17.2)\n",
      "Requirement already satisfied: matplotlib in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from mediapy) (3.8.2)\n",
      "Requirement already satisfied: absl-py>=0.7.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from dm_control>=1.0.3->D4RL==1.1) (2.1.0)\n",
      "Collecting dm-env (from dm_control>=1.0.3->D4RL==1.1)\n",
      "  Downloading dm_env-1.6-py3-none-any.whl.metadata (966 bytes)\n",
      "Collecting dm-tree!=0.1.2 (from dm_control>=1.0.3->D4RL==1.1)\n",
      "  Downloading dm_tree-0.1.9-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (2.2 kB)\n",
      "Requirement already satisfied: glfw in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from dm_control>=1.0.3->D4RL==1.1) (2.8.0)\n",
      "Collecting labmaze (from dm_control>=1.0.3->D4RL==1.1)\n",
      "  Downloading labmaze-1.0.6-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (278 bytes)\n",
      "Collecting lxml (from dm_control>=1.0.3->D4RL==1.1)\n",
      "  Downloading lxml-5.3.1-cp310-cp310-manylinux_2_28_x86_64.whl.metadata (3.7 kB)\n",
      "Collecting mujoco>=3.2.7 (from dm_control>=1.0.3->D4RL==1.1)\n",
      "  Downloading mujoco-3.2.7-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (44 kB)\n",
      "Requirement already satisfied: pyopengl>=3.1.4 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from dm_control>=1.0.3->D4RL==1.1) (3.1.9)\n",
      "Requirement already satisfied: pyparsing>=3.0.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from dm_control>=1.0.3->D4RL==1.1) (3.2.1)\n",
      "Requirement already satisfied: setuptools!=50.0.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from dm_control>=1.0.3->D4RL==1.1) (75.8.0)\n",
      "Requirement already satisfied: scipy in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from dm_control>=1.0.3->D4RL==1.1) (1.11.4)\n",
      "Requirement already satisfied: tqdm in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from dm_control>=1.0.3->D4RL==1.1) (4.67.1)\n",
      "Requirement already satisfied: decorator in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from ipython->mediapy) (5.1.1)\n",
      "Requirement already satisfied: jedi>=0.16 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from ipython->mediapy) (0.19.2)\n",
      "Requirement already satisfied: matplotlib-inline in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from ipython->mediapy) (0.1.7)\n",
      "Requirement already satisfied: prompt-toolkit!=3.0.37,<3.1.0,>=3.0.30 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from ipython->mediapy) (3.0.50)\n",
      "Requirement already satisfied: pygments>=2.4.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from ipython->mediapy) (2.19.1)\n",
      "Requirement already satisfied: stack-data in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from ipython->mediapy) (0.6.3)\n",
      "Requirement already satisfied: traitlets>=5 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from ipython->mediapy) (5.14.3)\n",
      "Requirement already satisfied: exceptiongroup in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from ipython->mediapy) (1.2.2)\n",
      "Requirement already satisfied: pexpect>4.3 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from ipython->mediapy) (4.9.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from jinja2->torch==2.0.1+cu117) (3.0.2)\n",
      "Requirement already satisfied: contourpy>=1.0.1 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from matplotlib->mediapy) (1.3.1)\n",
      "Requirement already satisfied: cycler>=0.10 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from matplotlib->mediapy) (0.12.1)\n",
      "Requirement already satisfied: fonttools>=4.22.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from matplotlib->mediapy) (4.55.8)\n",
      "Requirement already satisfied: kiwisolver>=1.3.1 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from matplotlib->mediapy) (1.4.8)\n",
      "Requirement already satisfied: packaging>=20.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from matplotlib->mediapy) (24.2)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from matplotlib->mediapy) (2.9.0.post0)\n",
      "Collecting Cython>=0.27.2 (from mujoco_py->D4RL==1.1)\n",
      "  Downloading Cython-3.0.12-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (3.3 kB)\n",
      "Requirement already satisfied: imageio>=2.1.2 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from mujoco_py->D4RL==1.1) (2.37.0)\n",
      "Requirement already satisfied: cffi>=1.10 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from mujoco_py->D4RL==1.1) (1.17.1)\n",
      "Collecting fasteners~=0.15 (from mujoco_py->D4RL==1.1)\n",
      "  Downloading fasteners-0.19-py3-none-any.whl.metadata (4.9 kB)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from requests->torchvision==0.15.2+cu117) (3.4.1)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from requests->torchvision==0.15.2+cu117) (3.10)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from requests->torchvision==0.15.2+cu117) (2.3.0)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from requests->torchvision==0.15.2+cu117) (2025.1.31)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from sympy->torch==2.0.1+cu117) (1.3.0)\n",
      "Requirement already satisfied: pycparser in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from cffi>=1.10->mujoco_py->D4RL==1.1) (2.22)\n",
      "Requirement already satisfied: attrs>=18.2.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from dm-tree!=0.1.2->dm_control>=1.0.3->D4RL==1.1) (25.1.0)\n",
      "Collecting wrapt>=1.11.2 (from dm-tree!=0.1.2->dm_control>=1.0.3->D4RL==1.1)\n",
      "  Downloading wrapt-1.17.2-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (6.4 kB)\n",
      "Requirement already satisfied: parso<0.9.0,>=0.8.4 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from jedi>=0.16->ipython->mediapy) (0.8.4)\n",
      "Requirement already satisfied: etils[epath] in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from mujoco>=3.2.7->dm_control>=1.0.3->D4RL==1.1) (1.12.0)\n",
      "Requirement already satisfied: ptyprocess>=0.5 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from pexpect>4.3->ipython->mediapy) (0.7.0)\n",
      "Requirement already satisfied: wcwidth in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from prompt-toolkit!=3.0.37,<3.1.0,>=3.0.30->ipython->mediapy) (0.2.13)\n",
      "Requirement already satisfied: six>=1.5 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from python-dateutil>=2.7->matplotlib->mediapy) (1.17.0)\n",
      "Requirement already satisfied: executing>=1.2.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from stack-data->ipython->mediapy) (2.2.0)\n",
      "Requirement already satisfied: asttokens>=2.1.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from stack-data->ipython->mediapy) (3.0.0)\n",
      "Requirement already satisfied: pure-eval in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from stack-data->ipython->mediapy) (0.2.3)\n",
      "Requirement already satisfied: fsspec in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from etils[epath]->mujoco>=3.2.7->dm_control>=1.0.3->D4RL==1.1) (2025.2.0)\n",
      "Requirement already satisfied: importlib_resources in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from etils[epath]->mujoco>=3.2.7->dm_control>=1.0.3->D4RL==1.1) (6.5.2)\n",
      "Requirement already satisfied: zipp in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from etils[epath]->mujoco>=3.2.7->dm_control>=1.0.3->D4RL==1.1) (3.21.0)\n",
      "Downloading protobuf-3.20.1-cp310-cp310-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (1.1 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.1/1.1 MB\u001b[0m \u001b[31m140.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading Pillow-9.0.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (4.3 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.3/4.3 MB\u001b[0m \u001b[31m175.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading triton-2.0.0-1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (63.3 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m63.3/63.3 MB\u001b[0m \u001b[31m202.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hDownloading einops-0.8.1-py3-none-any.whl (64 kB)\n",
      "Downloading mediapy-1.2.2-py3-none-any.whl (26 kB)\n",
      "Downloading dm_control-1.0.27-py3-none-any.whl (56.3 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m56.3/56.3 MB\u001b[0m \u001b[31m184.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hDownloading gym_notices-0.0.8-py3-none-any.whl (3.0 kB)\n",
      "Downloading h5py-3.12.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (5.3 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.3/5.3 MB\u001b[0m \u001b[31m152.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading mujoco_py-2.1.2.14-py3-none-any.whl (2.4 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.4/2.4 MB\u001b[0m \u001b[31m151.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading pybullet-3.2.7-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (103.2 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m103.2/103.2 MB\u001b[0m \u001b[31m150.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hDownloading Cython-3.0.12-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (3.6 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.6/3.6 MB\u001b[0m \u001b[31m181.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading dm_tree-0.1.9-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (152 kB)\n",
      "Downloading fasteners-0.19-py3-none-any.whl (18 kB)\n",
      "Downloading mujoco-3.2.7-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (6.4 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m6.4/6.4 MB\u001b[0m \u001b[31m173.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading cmake-3.31.4-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (27.8 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m27.8/27.8 MB\u001b[0m \u001b[31m221.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading dm_env-1.6-py3-none-any.whl (26 kB)\n",
      "Downloading labmaze-1.0.6-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (4.9 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.9/4.9 MB\u001b[0m \u001b[31m162.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading lit-18.1.8-py3-none-any.whl (96 kB)\n",
      "Downloading lxml-5.3.1-cp310-cp310-manylinux_2_28_x86_64.whl (5.2 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.2/5.2 MB\u001b[0m \u001b[31m198.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading wrapt-1.17.2-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (82 kB)\n",
      "Building wheels for collected packages: gym, D4RL, mjrl\n",
      "  Building wheel for gym (pyproject.toml) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for gym: filename=gym-0.23.1-py3-none-any.whl size=701426 sha256=15f82ac4bfe4147aaad6d0f913ff41b813a84085d2c64b88c0bed68bd1ccb7e1\n",
      "  Stored in directory: /home/zeus/.cache/pip/wheels/1a/00/fb/fe5cf2860fb9b7bc860e28f00095a1f42c7b726dd6f42d1acc\n",
      "  Building wheel for D4RL (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for D4RL: filename=D4RL-1.1-py3-none-any.whl size=26412345 sha256=55419958d2c8899fe20a7de17b03908cf493598d2fb33ea5bc743ac86f0407ab\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-cxisi6vd/wheels/7a/7a/27/d17500f4699272a90767c018dbd88a5e2376f2870a79b6a4ac\n",
      "  Building wheel for mjrl (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for mjrl: filename=mjrl-1.0.0-py3-none-any.whl size=61962 sha256=ec762cf0195d77fce72265c2be84581f7f11f0bb58fd2d7edd4988c28daa6e06\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-cxisi6vd/wheels/8f/99/f9/efd223b38d503df5eaada10ffe96a869fb0c0f3c92d9e43ed0\n",
      "Successfully built gym D4RL mjrl\n",
      "Installing collected packages: pybullet, mjrl, lit, gym_notices, wrapt, protobuf, Pillow, lxml, labmaze, h5py, gym, fasteners, einops, Cython, cmake, dm-tree, mujoco_py, mujoco, mediapy, dm-env, dm_control, D4RL, triton, torch, torchvision, torchaudio\n",
      "  Attempting uninstall: protobuf\n",
      "    Found existing installation: protobuf 4.23.4\n",
      "    Uninstalling protobuf-4.23.4:\n",
      "      Successfully uninstalled protobuf-4.23.4\n",
      "  Attempting uninstall: Pillow\n",
      "    Found existing installation: pillow 11.1.0\n",
      "    Uninstalling pillow-11.1.0:\n",
      "      Successfully uninstalled pillow-11.1.0\n",
      "  Attempting uninstall: mujoco\n",
      "    Found existing installation: mujoco 3.1.6\n",
      "    Uninstalling mujoco-3.1.6:\n",
      "      Successfully uninstalled mujoco-3.1.6\n",
      "  Attempting uninstall: triton\n",
      "    Found existing installation: triton 2.2.0\n",
      "    Uninstalling triton-2.2.0:\n",
      "      Successfully uninstalled triton-2.2.0\n",
      "  Attempting uninstall: torch\n",
      "    Found existing installation: torch 2.2.1+cu121\n",
      "    Uninstalling torch-2.2.1+cu121:\n",
      "      Successfully uninstalled torch-2.2.1+cu121\n",
      "  Attempting uninstall: torchvision\n",
      "    Found existing installation: torchvision 0.17.1+cu121\n",
      "    Uninstalling torchvision-0.17.1+cu121:\n",
      "      Successfully uninstalled torchvision-0.17.1+cu121\n",
      "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "gymnasium-robotics 1.3.1 requires mujoco<3.2.0,>=2.2.0, but you have mujoco 3.2.7 which is incompatible.\n",
      "lightning 2.5.0.post0 requires torch<4.0,>=2.1.0, but you have torch 2.0.1+cu117 which is incompatible.\n",
      "pytorch-lightning 2.5.0.post0 requires torch>=2.1.0, but you have torch 2.0.1+cu117 which is incompatible.\u001b[0m\u001b[31m\n",
      "\u001b[0mSuccessfully installed Cython-3.0.12 D4RL-1.1 Pillow-9.0.0 cmake-3.31.4 dm-env-1.6 dm-tree-0.1.9 dm_control-1.0.27 einops-0.8.1 fasteners-0.19 gym-0.23.1 gym_notices-0.0.8 h5py-3.12.1 labmaze-1.0.6 lit-18.1.8 lxml-5.3.1 mediapy-1.2.2 mjrl-1.0.0 mujoco-3.2.7 mujoco_py-2.1.2.14 protobuf-3.20.1 pybullet-3.2.7 torch-2.0.1+cu117 torchaudio-2.0.2+cu117 torchvision-0.15.2+cu117 triton-2.0.0 wrapt-1.17.2\n"
     ]
    }
   ],
   "source": [
    "!pip install torch==2.0.1+cu117 \\\n",
    "  torchvision==0.15.2+cu117 \\\n",
    "  torchaudio==2.0.2+cu117 \\\n",
    "  git+https://github.com/rail-berkeley/d4rl.git \\\n",
    "  gym==0.23.1 \\\n",
    "  protobuf==3.20.1 \\\n",
    "  einops \\\n",
    "  mediapy \\\n",
    "  Pillow==9.0.0 \\\n",
    "  -f https://download.pytorch.org/whl/torch_stable.html\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transferred weights for 0 layers\n",
      "Action shape: torch.Size([1, 16, 2])\n",
      "\n",
      "Predicted trajectory:\n",
      "Step  0: x=  48.9, y= 449.9\n",
      "Step  1: x= 118.0, y=  49.0\n",
      "Step  2: x= 191.8, y= 110.5\n",
      "Step  3: x= 501.7, y= 512.0\n",
      "Step  4: x=   0.0, y= 425.6\n",
      "Step  5: x= 378.3, y=   0.0\n",
      "Step  6: x=  39.3, y=   0.5\n",
      "Step  7: x= 474.6, y= 372.0\n",
      "Step  8: x=  17.0, y= 398.2\n",
      "Step  9: x=  30.2, y= 369.9\n",
      "Step 10: x=  11.4, y= 503.2\n",
      "Step 11: x= 512.0, y= 424.3\n",
      "Step 12: x= 415.7, y= 508.0\n",
      "Step 13: x= 357.8, y= 503.9\n",
      "Step 14: x= 294.6, y= 512.0\n",
      "Step 15: x= 219.7, y=  87.5\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from huggingface_hub import hf_hub_download\n",
    "from diffusers import DDPMScheduler\n",
    "\n",
    "class ObservationEncoder(nn.Module):\n",
    "\n",
    "    def __init__(self, state_dim):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(state_dim, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 256)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "class ObservationProjection(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.weight = nn.Parameter(torch.randn(32, 512))\n",
    "        self.bias = nn.Parameter(torch.zeros(32))\n",
    "\n",
    "    def forward(self, x):\n",
    "        if x.size(-1) == 256:\n",
    "            x = torch.cat([x, torch.zeros(*x.shape[:-1], 256, device=x.device)], dim=-1)\n",
    "        return nn.functional.linear(x, self.weight, self.bias)\n",
    "\n",
    "class UNet1D(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, hidden_channels=128):\n",
    "        super().__init__()\n",
    "        \n",
    "        # Downsampling path\n",
    "        self.down1 = nn.Sequential(\n",
    "            nn.Conv1d(in_channels, hidden_channels, kernel_size=3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv1d(hidden_channels, hidden_channels, kernel_size=3, padding=1),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "        \n",
    "        # Time embedding\n",
    "        self.time_mlp = nn.Sequential(\n",
    "            nn.Linear(1, hidden_channels),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_channels, hidden_channels)\n",
    "        )\n",
    "        \n",
    "        # Middle\n",
    "        self.mid = nn.Sequential(\n",
    "            nn.Conv1d(hidden_channels, hidden_channels, kernel_size=3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv1d(hidden_channels, hidden_channels, kernel_size=3, padding=1),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "        \n",
    "        # Upsampling path\n",
    "        self.up1 = nn.Sequential(\n",
    "            nn.Conv1d(2 * hidden_channels, hidden_channels, kernel_size=3, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv1d(hidden_channels, out_channels, kernel_size=3, padding=1)\n",
    "        )\n",
    "\n",
    "    def forward(self, x, t):\n",
    "        # Ensure proper tensor dimensions\n",
    "        if not isinstance(t, torch.Tensor):\n",
    "            t = torch.tensor([t], device=x.device)\n",
    "        if t.dim() == 0:\n",
    "            t = t.view(1)\n",
    "        if t.dim() == 1:\n",
    "            t = t.unsqueeze(-1)\n",
    "\n",
    "        # Time embedding\n",
    "        t_emb = self.time_mlp(t.float())  # [B, H]\n",
    "        \n",
    "        # Reshape time embedding to match spatial dimensions\n",
    "        t_emb = t_emb.unsqueeze(-1)  # [B, H, 1]\n",
    "        t_emb = t_emb.expand(-1, -1, x.shape[-1])  # [B, H, L]\n",
    "        \n",
    "        # Downsampling\n",
    "        d1 = self.down1(x)  # [B, H, L]\n",
    "        \n",
    "        # Add time embedding\n",
    "        mid = self.mid(d1 + t_emb)  # [B, H, L]\n",
    "        \n",
    "        # Upsampling with skip connections\n",
    "        up = self.up1(torch.cat([mid, d1], dim=1))  # [B, out_channels, L]\n",
    "        \n",
    "        return up\n",
    "\n",
    "class DiffusionPolicy:\n",
    "    def __init__(self, state_dim=5, device=\"cpu\"):\n",
    "        self.device = device\n",
    "        \n",
    "        # Define valid ranges\n",
    "        self.stats = {\n",
    "            \"obs\": {\n",
    "                \"min\": torch.zeros(5, device=device),\n",
    "                \"max\": torch.tensor([512, 512, 512, 512, 2 * np.pi], device=device)\n",
    "            },\n",
    "            \"action\": {\n",
    "                \"min\": torch.zeros(2, device=device),\n",
    "                \"max\": torch.full((2,), 512, device=device)\n",
    "            },\n",
    "        }\n",
    "        \n",
    "        self.obs_encoder = ObservationEncoder(state_dim).to(device)\n",
    "        self.obs_projection = ObservationProjection().to(device)\n",
    "        \n",
    "        # Use custom UNet1D implementation\n",
    "        self.model = UNet1D(\n",
    "            in_channels=34,  # 2 action channels + 32 context channels\n",
    "            out_channels=2,  # x,y coordinates\n",
    "            hidden_channels=128\n",
    "        ).to(device)\n",
    "        \n",
    "        self.noise_scheduler = DDPMScheduler(\n",
    "            num_train_timesteps=100,\n",
    "            beta_schedule=\"squaredcos_cap_v2\"\n",
    "        )\n",
    "        \n",
    "        # Load pre-trained weights using a more compatible approach\n",
    "        try:\n",
    "            checkpoint_path = hf_hub_download(\"dorsar/diffusion_policy\", \"push_tblock.pt\")\n",
    "            checkpoint = torch.load(checkpoint_path, map_location=device)\n",
    "            \n",
    "            # Load weights for encoder and projection\n",
    "            self.obs_encoder.load_state_dict(self._fix_state_dict(checkpoint[\"encoder_state_dict\"]))\n",
    "            self.obs_projection.load_state_dict(self._fix_state_dict(checkpoint[\"projection_state_dict\"]))\n",
    "            \n",
    "            # Transfer UNet weights\n",
    "            self._transfer_weights(checkpoint[\"model_state_dict\"])\n",
    "            \n",
    "        except Exception as e:\n",
    "            print(f\"Warning: Could not load pre-trained weights: {e}\")\n",
    "            print(\"The model will use randomly initialized weights.\")\n",
    "\n",
    "    def _fix_state_dict(self, state_dict):\n",
    "        \"\"\"Helper function to fix state dict keys if needed\"\"\"\n",
    "        new_state_dict = {}\n",
    "        for k, v in state_dict.items():\n",
    "            # Remove 'module.' prefix if it exists \n",
    "            k = k.replace('module.', '')\n",
    "            new_state_dict[k] = v\n",
    "        return new_state_dict\n",
    "\n",
    "    def _transfer_weights(self, original_state_dict):\n",
    "\n",
    "        custom_state_dict = self.model.state_dict()\n",
    "        \n",
    "        # Create mapping between original and custom architecture\n",
    "        layer_mapping = {\n",
    "            'down_blocks.0.resnets.0': 'down1.0',\n",
    "            'down_blocks.0.resnets.1': 'down1.2',\n",
    "            'mid_block.resnets.0': 'mid.0',\n",
    "            'mid_block.resnets.1': 'mid.2',\n",
    "            'up_blocks.0.resnets.0': 'up1.0',\n",
    "            'up_blocks.0.resnets.1': 'up1.2',\n",
    "        }\n",
    "        \n",
    "        # Transfer weights for compatible layers\n",
    "        transferred = set()\n",
    "        for orig_name, param in original_state_dict.items():\n",
    "            for orig_prefix, custom_prefix in layer_mapping.items():\n",
    "                if orig_name.startswith(orig_prefix):\n",
    "                    custom_name = orig_name.replace(orig_prefix, custom_prefix)\n",
    "                    if custom_name in custom_state_dict:\n",
    "                        if custom_state_dict[custom_name].shape == param.shape:\n",
    "                            custom_state_dict[custom_name].copy_(param)\n",
    "                            transferred.add(custom_name)\n",
    "        \n",
    "        # Load the transferred weights\n",
    "        self.model.load_state_dict(custom_state_dict, strict=False)\n",
    "        \n",
    "        print(f\"Transferred weights for {len(transferred)} layers\")\n",
    "\n",
    "    def normalize_data(self, data, stats):\n",
    "        return ((data - stats[\"min\"]) / (stats[\"max\"] - stats[\"min\"])) * 2 - 1\n",
    "\n",
    "    def unnormalize_data(self, ndata, stats):\n",
    "        return ((ndata + 1) / 2) * (stats[\"max\"] - stats[\"min\"]) + stats[\"min\"]\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def predict(self, observation):\n",
    "        # Ensure observation is a tensor and has batch dimension\n",
    "        if not isinstance(observation, torch.Tensor):\n",
    "            observation = torch.tensor(observation, device=self.device)\n",
    "        if observation.dim() == 1:\n",
    "            observation = observation.unsqueeze(0)\n",
    "            \n",
    "        observation = observation.to(self.device)\n",
    "        normalized_obs = self.normalize_data(observation, self.stats[\"obs\"])\n",
    "        \n",
    "        # Generate context\n",
    "        cond = self.obs_projection(self.obs_encoder(normalized_obs))\n",
    "        cond = cond.view(normalized_obs.shape[0], -1, 1).expand(-1, -1, 16)\n",
    "        \n",
    "        # Initialize with noise\n",
    "        action = torch.randn((observation.shape[0], 2, 16), device=self.device)\n",
    "        \n",
    "        # Denoise\n",
    "        self.noise_scheduler.set_timesteps(100)\n",
    "        for t in self.noise_scheduler.timesteps:\n",
    "            model_input = torch.cat([action, cond], dim=1)\n",
    "            model_output = self.model(model_input, t.to(self.device))\n",
    "            action = self.noise_scheduler.step(model_output, t.to(self.device), action).prev_sample\n",
    "        \n",
    "        action = action.transpose(1, 2)\n",
    "        action = self.unnormalize_data(action, self.stats[\"action\"])\n",
    "        return action\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    policy = DiffusionPolicy()\n",
    "    \n",
    "    # Test with sample observation\n",
    "    obs = torch.tensor([[\n",
    "        256.0,  # robot arm x position\n",
    "        256.0,  # robot arm y position\n",
    "        200.0,  # block x position\n",
    "        300.0,  # block y position\n",
    "        np.pi / 2,  # block angle\n",
    "    ]])\n",
    "    \n",
    "    action = policy.predict(obs)\n",
    "    print(\"Action shape:\", action.shape)\n",
    "    print(\"\\nPredicted trajectory:\")\n",
    "    for i, (x, y) in enumerate(action[0]):\n",
    "        print(f\"Step {i:2d}: x={x:6.1f}, y={y:6.1f}\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
